package com.ztom.agent.asm;

import com.ztom.agent.asm.plugin.PluginClassVisitor;
import com.ztom.agent.asm.plugin.PluginInfoHolder;
import com.ztom.agent.config.TraceConfig;
import com.ztom.agent.config.TraceFilter;
import com.ztom.agent.log.Logger;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.ClassWriter;

import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.nio.file.StandardOpenOption;
import java.security.ProtectionDomain;

/**
 * @author ZhangTao
 */
public class LinkTraceTransformer implements ClassFileTransformer {
    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {

        try {
            // 略过在[不需要注入]的类
            if (TraceFilter.isNotNeedInject(className)) {
                return classfileBuffer;
            }

            // 略过不在[需要注入]中且不是[插件注入]的类
            if (!TraceFilter.isNeedInject(className) && !PluginInfoHolder.getClassSet().contains(className)) {
                return classfileBuffer;
            }
            Logger.debug("LinkTraceTransformer.transform(" + getClassLoaderName(loader) + ", " + className + ", classBeingRedefined, protectionDomain, " + classfileBuffer.length + ")...");
            byte[] bytes = getBytes(className, classfileBuffer);

            saveUpdatedClass(className, bytes);
            return bytes;
        } catch (Exception e) {
            Logger.error("LinkTraceTransformer.transform(" + getClassLoaderName(loader) + ", " + className + ", " + classBeingRedefined + ", protectionDomain, " + classfileBuffer.length + ")", e);
        }
        return classfileBuffer;
    }

    private byte[] getBytes(String className, byte[] classFileBuffer) {
        ClassReader cr = new ClassReader(classFileBuffer);
        ClassWriter cw = new ClassWriter(cr, ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS);

        ClassVisitor cv;
        if (PluginInfoHolder.getClassSet().contains(className)) {
            cv = new PluginClassVisitor(cw, className);
        } else {
            cv = new TraceClassVisitor(cw, className);
        }

        cr.accept(cv, ClassReader.SKIP_FRAMES | ClassReader.SKIP_DEBUG);
        return cw.toByteArray();
    }

    private void saveUpdatedClass(String className, byte[] bytes) throws IOException {
        // 保存修改后的class文件
        if (TraceConfig.getInstance().isDumpClassFile()) {
            String rootPath = TraceConfig.getInstance().getDumpClassFilePath();
            if (!Files.exists(Paths.get(rootPath))) {
                Files.createDirectories(Paths.get(rootPath));
            }
            String dumpFilePath = rootPath + className.replace("/", "_") + ".class";
            Files.write(Paths.get(dumpFilePath), bytes,
                    StandardOpenOption.WRITE, StandardOpenOption.CREATE, StandardOpenOption.TRUNCATE_EXISTING);
            Logger.debug("LinkTraceTransformer.transform dump the updated class file, file path=" + dumpFilePath);
        }
    }

    private String getClassLoaderName(ClassLoader classLoader) {
        if (classLoader == null) {
            return "null";
        }

        return classLoader.getClass().getName();
    }
}
